# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import getpass
import logging
import subprocess
import sys
import tempfile
import os
import shutil

from absl import app, flags
import jinja2

from coin.base.datetime_util import iterate_date
from coin.support.feed_tool.feed_cache.app.coin2_fast_feed_runner import (load_feed_writer_config,
                                                                          load_subscribers_setting,
                                                                          get_mea_from_config,
                                                                          get_recipe_from_config,
                                                                          get_groups_from_setting)

FLAGS = flags.FLAGS


class NoJobError(Exception):
  pass


temp_driver_files = []
def generate_script_setting(trading_dates):
  machines = [
      "feed-02.ap-northeast-2.aws",
      "feed-05.ap-northeast-1.aws",
      "feed-05.cn-hongkong.aliyun",
      "feed-02.eu-west-1.aws",
      "feed-02.us-east-1.aws",
      "feed-01.ap-northeast-1.aws.huobi",
      "feed-07.ap-northeast-1.aws",
  ]
  ignore_meas = {
      'Futures.Okex.v3-swap',
      'Spot.Okex.v3',
  }
  config, config_file = load_feed_writer_config(machines)
  subscribers = load_subscribers_setting()
  settings = {}
  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)
  template = template_env.get_template(
      'coin/support/feed_tool/feed_cache/app/index_generator/archive_config.tmpl')
  tmp_dir = '/remote/iosg/jenkins-2/bot-coin/jenkins/workspace/coin2_fastfeed/coin/tmp'
  for trading_date in trading_dates:
    settings[trading_date] = {}
    for m in machines:
      if (FLAGS.machine is not None) and (m not in FLAGS.machine):
        continue
      settings[trading_date][m] = {}
      meas = get_mea_from_config(m, config)
      for mea in meas:
        if (FLAGS.mea is not None) and (mea not in FLAGS.mea):
          continue
        if mea in ignore_meas:
          continue
        #target_mear = '{}.index'.format(mea)
        #target_dir = os.path.join('/remote/iosg/coin-2/feed_cache2/', m, trading_date, target_mear)
        #print('Delete targe_dir first: {}'.format(target_dir))
        #shutil.rmtree(target_dir, ignore_errors=True)
        settings[trading_date][m][mea] = {}
        recipes = get_recipe_from_config(m, mea, config, subscribers)
        groups = get_groups_from_setting(m, mea, config)
        for recipe in recipes:
          if 'non_tbs' not in recipe:
            continue
          if (FLAGS.recipe is not None) and (recipe not in FLAGS.recipe):
            continue
          settings[trading_date][m][mea][recipe] = {}
          for group in groups:
            if (FLAGS.group is not None) and (group not in FLAGS.group):
              continue
            driver_file = tempfile.NamedTemporaryFile(mode='w', dir=tmp_dir)
            driver = template.render(
                machine=m,
                start=trading_date,
                mea=mea,
                group=group,
                recipe=recipe)
            driver_file.write(driver)
            driver_file.flush()
            temp_driver_files.append(driver_file)
            settings[trading_date][m][mea][recipe][group] = driver_file.name
  return settings


def generate_slurm_script(
    start_date,
    end_date,
    git_commit,
    summary_mailto,
    user=None):
  start_date = datetime.datetime.strptime(start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')
  trading_dates = [td for td in iterate_date(start_date.date(), end_date.date())]
  trading_dates = [td.strftime('%Y%m%d') for td in trading_dates]

  settings = generate_script_setting(trading_dates)

  mail_subject = "[index-generator] ${SLURM_JOB_ID} summary"

  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)
  template = template_env.get_template(
      'coin/support/feed_tool/feed_cache/app/index_generator/slurm_script.tmpl')
  return template.render(settings=settings,
                         mail_subject=mail_subject,
                         summary_mailto=summary_mailto)


class SBatchError(Exception):
  pass


def run_sbatch(script_filepath, sbatch_args=None, verbose=False):
  sbatch_args = sbatch_args or []
  if isinstance(sbatch_args, str):
    sbatch_args = [token for token in sbatch_args.split(' ') if token]

  sbatch_args = ['--parsable'] + sbatch_args
  if sbatch_args:
    logging.debug('sbatch_args: %s', ' '.join(sbatch_args))

  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath], stdout=subprocess.PIPE)
  if completed.returncode != 0:
    raise SBatchError(completed)

  sbatch_output = completed.stdout.decode('utf-8')
  job_id = int(sbatch_output.strip())
  if verbose:
    sys.stderr.write('Job ID: %d\n' % job_id)
  return job_id


def run_on_slurm(
    start_date,
    end_date,
    git_commit,
    summary_mailto='{user}@prestolabs.io',
    dry_run=False,
    sbatch_args=None,
    verbose=False):
  if summary_mailto is not None:
    if '{user}' in summary_mailto:
      user = 'leon'
      summary_mailto = summary_mailto.format(user=user)
    if verbose:
      sys.stderr.write('Index generation result mail will be sent to %s.\n' % summary_mailto)

  try:
    script = generate_slurm_script(start_date,
                                   end_date,
                                   git_commit,
                                   summary_mailto,
                                   user=getpass.getuser())
  except NoJobError:
    if verbose:
      sys.stderr.write('No job to run.\n')
    raise

  if dry_run:
    if verbose:
      print(script)
    return

  tmp_dir = '/remote/iosg/jenkins-2/bot-coin/jenkins/workspace/coin2_fastfeed/coin/tmp'
  with tempfile.NamedTemporaryFile(mode='w', dir=tmp_dir) as script_file:
    script_file.write(script)
    script_file.flush()

    try:
      job_id = run_sbatch(script_file.name, sbatch_args, verbose=verbose)
      if verbose:
        print(job_id)
    except SBatchError:
      if verbose:
        sys.stderr.write('Error while running sbatch. Abort.\n')
      raise

  for tmp_driver in temp_driver_files:
    tmp_driver.close()
  return job_id


def main(argv):
  try:
    run_on_slurm(start_date=FLAGS.start_date,
                 end_date=FLAGS.end_date,
                 git_commit=FLAGS.git_commit,
                 summary_mailto=FLAGS.mailto,
                 dry_run=FLAGS.script_only,
                 sbatch_args=FLAGS.sbatch_args,
                 verbose=True)
  except NoJobError:
    return 0
  except SBatchError:
    return 1


if __name__ == '__main__':
  flags.DEFINE_string('start_date', None, 'yyyymmdd')
  flags.DEFINE_string('end_date', None, 'yyyymmdd. Feed at end_date is not checked.')
  flags.DEFINE_string('git_commit', 'HEAD', 'Git commit sha or HEAD.')
  flags.DEFINE_string('mailto',
                      '{user}@prestolabs.io',
                      'If specified, it will send a mail after finish.')
  flags.DEFINE_bool('script_only',
                    False,
                    'If true, it will print sbatch script and not run sbatch.')
  flags.DEFINE_string(
      'sbatch_args',
      '-N10 --ntasks=530 --mem-per-cpu=1G --job-name=index --priority=TOP --wait',
      'sbatch arguments')

  flags.DEFINE_list('machine', None, "Separate by comma")

  flags.DEFINE_list('mea', None, "Separate by comma")

  flags.DEFINE_list('recipe', None, "Separate by comma")

  flags.DEFINE_list('group', None, "Separate by comma")

  app.run(main)
